#!/bin/bash

#
# skssh -- by Dario Berzano <dario.berzano@gmail.com>
#
# Part of the lfscripts suite -- http://lfscripts.googecode.com/
#
# Automatically choose the proper ssh connection for shell/rsync for the current
# network topology. It remembers the last selection, so changing network
# topology only introduces a little delay the first time.
#
# This tool can do: ssh, scp and rsync (in both directions). Do a symbolic link
# to the main skssh and give it proper names to trigger different configs.
#

#
# Variables
#

# Settings prefix (it is a Bash file to source)
export SETTINGS_PREFIX="$HOME/.lfscripts/skssh_"

# Where to store the last working server (set in LoadConfig)
export PREVFILE

# Last working server, read from file (set in LoadConfig)
export PREVSRV

# SSH error code
export SSHERR=255

#
# Functions
#

# Smart echo
function Echo() {
  echo -e "$@" >&2
}

# Connect via ssh, rsync, scp: specify it via the first parameter (which may be
# "sh", "sync" or "cp"), and specify host[:port] as second parameter. Other
# parameters, if given, are passed to ssh/scp in "sh"/"cp" mode
function Connect() {

  local ACT=$1
  local SRV=$2
  shift 2

  Echo "\033[1;33mConnecting to \033[1;35m$SRV\033[1;33m...\033[m"

  # Separate host and (eventually) port
  local SSHHOST
  local SSHPORT
  SSHHOST=$( echo $SRV | perl -ne '/^([^:]+)(:([0-9]+))?/; print "$1"' )
  SSHPORT=$( echo $SRV | perl -ne '/^([^:]+)(:([0-9]+))?/; print "$3"' )
  [ "$SSHPORT" == "" ] && SSHPORT=22

  # Connect string
  local SSHCMD="ssh -p$SSHPORT $SSHHOST $SSHOPT"
  local RET

  if [ $ACT == "sh" ]; then

    # This is the shell case; check if we want a persistent shell
    local SCRN
    if [ "${1:0:2}" == "-s" ]; then
      expr "${1:2}" + 0 > /dev/null 2>&1
      if [ $? == 0 ]; then
        SCRN=${1:2}
      fi
    fi

    if [ "$SCRN" == "" ]; then
      $SSHCMD "$@"
      RET=$?
    else
      $SSHCMD -t screen -S scr$SCRN -d -RR
      RET=$?
    fi

  elif [ $ACT == "sync" ]; then

    # This is the rsync case
    $SSHCMD echo OK > /dev/null 2>&1
    RET=$?
    if [ $RET != $SSHERR ]; then

      # Prepare --include switches
      local RSYNCINC_SWITCH
      local L=${#RSYNCINC[@]}
      local I
      local J=0
      for ((I=0; $I<$L; I++)); do
        RSYNCINC_SWITCH[$J]='--include=**/'"${RSYNCINC[$I]}"
        let J++
        RSYNCINC_SWITCH[$J]='--include=**/'"${RSYNCINC[$I]}"'/**'
        let J++
      done

      # If it works: rsync; just decide the direction
      if [ "$1" == "--to-remote" ]; then
        Echo "Copying modified data to remote host..."
        rsync -a --delete --no-owner --no-group \
          "${RSYNCINC_SWITCH[@]}" --exclude='*' \
          -e "ssh -p$SSHPORT $SSHOPT" $RSYNCEXTRA \
          "$RSYNCSRC" $SSHHOST:"$RSYNCDST"
        RET=$?
      elif [ "$1" == "--from-remote" ]; then
        Echo "Copying modified data from remote host..."
        rsync -a --delete --no-owner --no-group \
          "${RSYNCINC_SWITCH[@]}" --exclude='*' \
          -e "ssh -p$SSHPORT $SSHOPT" $RSYNCEXTRA \
          $SSHHOST:"$RSYNCDST" "$RSYNCSRC"
        RET=$?
      else
        Echo "You have to explicitly specify \033[1;36m--to-remote\033[m or \033[1;36m--from-remote\033[m"
        RET=1
      fi

      if [ $RET == 0 ]; then
        Echo "\033[1;33mSync \033[1;32mOK\033[m"
      else
        Echo "\033[1;33mSync \033[1;31mFAILED\033[m"
      fi
    fi

  elif [ $ACT == "cp" ]; then

    local FSRC FDST

    # This is the scp case: just decide the direction
    if [ $# -lt 3 ]; then
      RET=1
    else
      if [ "$1" == "--to-remote" ]; then
        FSRC="$2"
        FDST="$SSHHOST:$3"
        RET=0
      elif [ "$1" == "--from-remote" ]; then
        FSRC="$SSHHOST:$2"
        FDST="$3"
        RET=0
      else
        RET=1
      fi
    fi

    if [ $RET == 0 ]; then
      shift 3
      scp $SSHOPT "$@" -oPort=$SSHPORT "$FSRC" "$FDST"
      RET=$?
      if [ $RET == 0 ]; then
        Echo "\033[1;33mCopy \033[1;32mOK\033[m"
      else
        Echo "\033[1;33mCopy \033[1;31mFAILED\033[m (with error=$RET)"
        RET=$SSHERR
      fi
    else
      Echo "Usage: $0 \033[1;36m--to-remote\033[m|\033[1;36m--from-remote\033[m <source_filename> <dest_filename>"
      RET=1
    fi

  fi

  # Write preference in file, if connection succeeded
  if [ $RET != $SSHERR ]; then
    echo $SRV > "$PREVFILE"
  fi

  return $RET
}

# Tries to load the configuration from the specified ($1) tag
function LoadConfig() {

  local CONFIG_FILE=${SETTINGS_PREFIX}$1

  if [ -r "$CONFIG_FILE" ]; then
    source "$CONFIG_FILE"

    # PREVFILE, PREVSRV are globals
    PREVFILE="$HOME/.$1sh_last"
    PREVSRV=$(cat "$PREVFILE" 2> /dev/null)

    return 0
  fi

  return 1
}

# Main function
function Main() {

  local CONFIG_TAG
  local BASE=$(basename "$0")

  if [[ "$BASE" =~ ^(.*)sh$ ]]; then
    CONFIG_TAG="${BASH_REMATCH[1]}"
    ACTION=sh
  elif [[ "$BASE" =~ ^(.*)cp$ ]]; then
    CONFIG_TAG="${BASH_REMATCH[1]}"
    ACTION=cp
  elif [[ "$BASE" =~ ^(.*)sync$ ]]; then
    CONFIG_TAG="${BASH_REMATCH[1]}"
    ACTION=sync
  else
    Echo "Error: executable suffix is not one of: sh, cp, sync"
    exit 1
  fi

  # Loads config
  LoadConfig "$CONFIG_TAG"
  if [ $? != 0 ]; then
    Echo "Can't load configuration file"
    exit 2
  fi

  # Possible values: 0:nomatch, 1:tried_and_failed
  local STATUS=0
  local RET

  # Search if last used server is in the list
  for SRV in $SERVERS
  do
    if [ "$PREVSRV" == "$SRV" ]; then
      Connect $ACTION $SRV "$@"
      RET=$?
      if [ $RET != $SSHERR ]; then
        # Connection ok: exit
        exit $RET
      else
        STATUS=1
        break
      fi
    fi
  done

  # If error occured (or no server found), try the list
  for SRV in $SERVERS
  do
    if [ $STATUS == 1 ] && [ "$PREVSRV" == "$SRV" ]; then
      continue
    fi
    Connect $ACTION $SRV "$@"
    RET=$?
    [ $RET != $SSHERR ] && exit $RET
  done

  # If this point is reached, every attempt failed
  Echo "\033[1;33mAll tries have \033[1;31mFAILED\033[1;33m, sorry\033[m" >&2
  rm -f "$PREVFILE"
  exit $SSHERR
}

#
# Entry point
#

Main "$@"
